feat(mla): expose dense MLA split precision options for Kimi-K3 - #595

• vllm/distributed/communication_op.py#L29-L36 (this comment)
• vllm/model_executor/models/kimi_k25_vit.py#L285-L306
• vllm/model_executor/models/kimi_k25_vit.py#L843-L860

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

🌐 Web query:

What architecture strings does torch.cuda.get_arch_list() return for Blackwell GPUs, and are CUDA cubins compatible across minor compute-capability versions?

💡 Result:

For NVIDIA Blackwell GPUs, torch.cuda.get_arch_list typically returns strings such as 'sm_100' or 'sm_120' [1], and potentially variants like 'sm_120a' when specifically configured to include accelerated features [2]. The exact list depends on your specific PyTorch version, installed CUDA toolkit version, and hardware [1]. Regarding CUDA binary compatibility, cubin files are compatible across minor compute-capability versions within the same major version, provided the target GPU has a minor version greater than or equal to the minor version for which the cubin was compiled [3][4]. Specifically: - Forward Compatibility: A cubin generated for compute capability X.Y is supported on any GPU with compute capability X.Z where Z >= Y [3][4]. - Backward Compatibility: A cubin is not supported on a GPU with a lower minor version (e.g., code compiled for 10.1 cannot run on 10.0) [3][4]. - Major Versions: Cubins are not binary compatible across different major compute-capability versions (e.g., code compiled for 10.x will not run on 12.x) [4]. To ensure broader compatibility and future-proofing, it is recommended to include PTX (Parallel Thread Execution) code in your binaries, as PTX is forward-compatible across different major and minor versions [3][4].

Citations:

• 1: GitHub issue 157549 in pytorch/pytorch (link omitted to avoid creating a cross-reference)
• 2: GitHub issue 172807 in pytorch/pytorch (link omitted to avoid creating a cross-reference)
• 3: https://docs.nvidia.com/cuda/blackwell-compatibility-guide/index.html
• 4: https://docs.nvidia.com/cuda/cuda-programming-guide/01-introduction/cuda-platform.html

🏁 Script executed:

# Inspect the repository conventions, scoped guidance, and the target function around
# the reported check. Also identify nearby tests or callers that define its contract.
printf '%s\n' '--- scoped repository guidance ---'
find /tmp/coderabbit-repo-knowledge/local-inference-lab-vllm-a28dce09 \
  -maxdepth 2 -type f -name '*.md' -print
printf '%s\n' '--- target file outline ---'
ast-grep outline vllm/entrypoints/k3_dspark_standalone.py
printf '%s\n' '--- target context ---'
cat -n vllm/entrypoints/k3_dspark_standalone.py | sed -n '130,220p'
printf '%s\n' '--- direct references ---'
rg -n --glob '*.py' 'get_device_capability|get_arch_list|expected_arch|k3_dspark_standalone' .

Repository: local-inference-lab/vllm

Length of output: 24592

🏁 Script executed:

#!/bin/bash
printf '%s\n' '--- repository-wide conventions ---'
cat /tmp/coderabbit-repo-knowledge/local-inference-lab-vllm-a28dce09/conventions/repo-wide.md
printf '%s\n' '--- scoped test guidance ---'
cat /tmp/coderabbit-repo-knowledge/local-inference-lab-vllm-a28dce09/learnings/tests.md
printf '%s\n' '--- standalone tests ---'
cat -n tests/v1/spec_decode/test_k3_dspark_standalone.py
printf '%s\n' '--- runtime loading call path ---'
cat -n vllm/entrypoints/k3_dspark_standalone.py | sed -n '360,430p'
printf '%s\n' '--- PyTorch arch-list references in repository ---'
rg -n --glob '*.py' 'get_arch_list|arch_list|compiled arches|compiled_arch' vllm tests

Repository: local-inference-lab/vllm

Length of output: 12488

🌐 Web query:

site:docs.nvidia.com CUDA sm_120a sm_120 architecture suffix compatibility cubin minor compute capability

💡 Result:

In the CUDA ecosystem, the sm_120a designation is a specific target architecture variant, distinct from the base sm_120 architecture [1][2]. Key technical details regarding these identifiers include: 1. Architecture Suffixes: The a suffix (as in sm_120a) denotes an architecture-specific variant [2]. These variants often include specialized hardware features or capacity limits that differ from the base architecture [3]. For example, sm_120a is documented with a specific maximum shared memory capacity (100 KB per CTA), which may differ from other variants in the same series [3]. 2. Compatibility and Compilation: - Code compiled for a specific sm_XYa target is generally restricted to that target and cannot be compiled to general sm_XY hardware [4]. - In contrast, PTX code compiled for a virtual architecture (e.g., compute_120) can be JIT-compiled to real binary targets (sm_*) of the same or higher compute capability [5]. - When specifying sm_120a as the target in nvcc, the compiler handles the inclusion of both architecture-specific and non-specific code as needed to maintain compatibility [2]. 3. Cubin and Binary Utilities: A cubin (CUDA binary) is a compiled device code file targeted at a specific SM version [5][6]. Tools like nvdisasm treat sm_120 and sm_120a as distinct, supported architectures [1][7]. When using binary utilities, the specified architecture must match the asserted architecture of the binary for accurate processing [1]. In summary, sm_120a is a specialized architecture target that provides unique features or constraints compared to the base sm_120. Developers should explicitly target these variants using the appropriate -arch or --gpu-architecture flags in the CUDA compiler driver to ensure the generated binaries are compatible with the intended hardware [2][7].

Citations:

• 1: https://docs.nvidia.com/cuda/archive/13.2.2/cuda-binary-utilities/index.html
• 2: https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html
• 3: https://docs.nvidia.com/cuda/parallel-thread-execution/index.html
• 4: https://docs.nvidia.com/cuda/archive/13.2.0/ptx-compiler-api/index.html
• 5: https://docs.nvidia.com/cuda/cuda-programming-guide/01-introduction/cuda-platform.html
• 6: https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/
• 7: https://docs.nvidia.com/cuda/cuda-binary-utilities/index.html

🌐 Web query:

site:pytorch.org torch.cuda.get_arch_list documentation compiled CUDA architectures

💡 Result:

The function torch.cuda.get_arch_list is a utility in the PyTorch library that returns a list of strings representing the CUDA architectures for which the installed version of PyTorch was compiled [1][2][3]. This function provides insight into the hardware compatibility of your specific PyTorch installation by identifying the compute capabilities (e.g., "sm_80", "sm_86") that were enabled during the compilation process [1][3]. Key details: - Return type: list[str] [1][2][3]. - Purpose: It allows users to verify whether their PyTorch binaries include support for the specific CUDA architecture of the GPU(s) they are using [1][3].

Citations:

• 1: https://docs.pytorch.org/docs/2.13/generated/torch.cuda.get%5Farch%5Flist.html
• 2: https://docs.pytorch.org/docs/stable/generated/torch.cuda.get_arch_list.md
• 3: https://docs.pytorch.org/docs/2.3/generated/torch.cuda.get_arch_list.html

🌐 Web query:

b12x attention dense_mla bind query_cache_seqlens parameter

💡 Result:

In the b12x framework, dense_mla is part of a library designed for high-performance attention mechanisms, including Multi-Head Latent Attention (MLA) [1][2]. The parameter query_cache_seqlens (often associated with metadata classes like MLASparseDecodeMetadata or MLASparseExtendMetadata) is used to specify the sequence lengths of the active tokens in the Key-Value (KV) cache during the decoding or extension phase [3][1][4]. In the b12x workflow, binding involves associating these metadata parameters with a workspace object that manages physical memory and ensures compatibility with CUDA Graph execution [1][5]. For example, when using the sparse MLA decode path, the workspace.prepare_decode method is called to bind metadata such as page_table_1 and cache_seqlens_int32 (the internal representation of query_cache_seqlens) to the internal workspace buffers [3][1]. This process allows the system to maintain stable metadata, which is critical for replaying kernels under CUDA Graphs without corruption [1][6][5]. The b12x library uses these bindings to manage complex memory layouts for latent KV caches, supporting different operational modes like decode and extend [1]. These metadata structures must be explicitly prepared and bound before the execution of the sparse_mla_decode_forward or sparse_mla_extend_forward functions, which ultimately trigger the computation [3][4]. It is important to distinguish this b12x implementation from other libraries like FlashMLA (developed by DeepSeek/ByteDance), which use different interfaces such as get_mla_metadata and flash_mla_with_kvcache [7][8][9][10]. While both serve similar purposes in optimizing MLA, the specific parameter names and binding contracts described (e.g., B12XAttentionWorkspace, MLASparseDecodeMetadata) are unique to the b12x API [1][4].

Citations:

• 1: https://deepwiki.com/bhaktatejas922/b12x/3.2-mla-workspace-and-api
• 2: https://github.com/local-inference-lab/b12x
• 3: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/b12x/attention/mla/api.py
• 4: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/tests/nsa_mla_numerical_reference_battery.py
• 5: https://deepwiki.com/bhaktatejas922/b12x/2-paged-attention-system
• 6: https://deepwiki.com/bhaktatejas922/b12x/9-testing-infrastructure
• 7: https://github.com/deepseek-ai/FlashMLA/blob/main/README.md
• 8: https://github.com/deepseek-ai/FlashMLA/blob/main/flash_mla/flash_mla_interface.py
• 9: https://github.com/bytedance-iaas/FlashMLA
• 10: https://github.com/PFCCLab/FlashMLA

🏁 Script executed:

#!/bin/bash
set -eu
printf '%s\n' '--- repository guidance ---'
find /tmp/coderabbit-repo-knowledge/local-inference-lab-vllm-a28dce09 -maxdepth 2 -type f -name '*.md' -print
printf '%s\n' '--- target file structure ---'
ast-grep outline vllm/v1/attention/backends/mla/b12x_mla.py
printf '%s\n' '--- target call and nearby imports ---'
sed -n '1,90p' vllm/v1/attention/backends/mla/b12x_mla.py
sed -n '940,1010p' vllm/v1/attention/backends/mla/b12x_mla.py
printf '%s\n' '--- local b12x references ---'
rg -n --glob '*.py' 'dense_mla|query_cache_seqlens|query_cache_seq_lens|bind\(' vllm tests 2>/dev/null | head -250

Repository: local-inference-lab/vllm

Length of output: 32548

🏁 Script executed:

#!/bin/bash
set -eu
printf '%s\n' '--- applicable repository conventions ---'
cat /tmp/coderabbit-repo-knowledge/local-inference-lab-vllm-a28dce09/conventions/repo-wide.md
printf '%s\n' '--- b12x dependency declarations ---'
rg -n -i --glob '!*lock*' --glob '!*.pyc' 'b12x|dense_mla' pyproject.toml setup.py setup.cfg requirements* docker .github vllm tests 2>/dev/null | head -250
printf '%s\n' '--- metadata construction and call context ---'
sed -n '170,275p' vllm/v1/attention/backends/mla/b12x_mla.py
sed -n '285,620p' vllm/v1/attention/backends/mla/b12x_mla.py
sed -n '730,1010p' vllm/v1/attention/backends/mla/b12x_mla.py
printf '%s\n' '--- relevant tests around bind keyword shape ---'
sed -n '200,235p' tests/v1/attention/test_b12x_mla.py
sed -n '630,675p' tests/v1/attention/test_b12x_mla.py

Repository: local-inference-lab/vllm

Length of output: 50380

🌐 Web query:

site:github.com/bhaktatejas922/b12x "def bind" "query_cache_seqlens"

💡 Result:

In the context of the b12x codebase, cache_seqlens is a tensor (specifically a [batch] tensor of type int32) representing the current sequence lengths stored in the KV cache for each request in a batch [1][2][3]. While there is no single function literally named bind_query_cache_seqlens, the term likely refers to the process of binding or preparing these sequence length tensors within the model's forward pass or workspace [4][3][5]. The workflow involves several key operations: 1. Build: The KVCache manager generates the cache_seqlens tensor, which tracks the number of tokens currently residing in the KV cache for each request [2][6]. 2. State Binding: During the forward pass, this tensor is included in a StepState object, which is passed to layers to ensure they have the necessary context for operations like paged attention [3]. 3. Workspace Preparation: Before execution, a workspace (such as PagedAttentionWorkspace) consumes these metadata tensors (including cache_seqlens, page_table, and cu_seqlens_q) to prepare for the attention computation [4][5]. Essentially, binding cache_seqlens is part of linking the engine's scheduler state to the actual low-level PagedAttention kernels [4][6][5].

Citations:

• 1: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/serve/engine/scheduler.py
• 2: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/serve/cache/kv_cache.py
• 3: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/serve/engine/step_state.py
• 4: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/serve/model/attention.py
• 5: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/b12x/attention/paged/api.py
• 6: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/serve/engine/runner.py

🌐 Web query:

site:github.com/local-inference-lab/b12x "dense_mla" "query_cache_seqlens"

💡 Result:

Dense MLA and query_cache_seqlens are technical components associated with high-performance large language model (LLM) inference libraries, specifically found in projects like b12x [1][2]. Dense MLA refers to Multi-Head Latent Attention (MLA) mechanisms optimized for dense computation, often involving fused kernel implementations to improve performance during model execution [1]. In these systems, MLA is used to manage attention operations, with specific kernels handling query projection and attention decoding [1]. Query_cache_seqlens is a parameter or data structure used within the inference runtime to manage sequence lengths for queries stored in the KV cache [1]. It is typically utilized during the binding and execution phases of stateful attention kernels (such as compressed or sparse MLA), where the runtime must track the specific lengths of sequences across different cache pages to ensure correct attention computation during the decode or extend phases [1]. In these architectures, the planning and runtime flow for such operations generally follows a pattern of: 1. Planning the work (e.g., compressed_mla.plan) [1]. 2. Sizing and allocating scratch space based on the plan [1]. 3. Binding tensors (using structures like query_cache_seqlens) as views [1]. 4. Running the CUDA-graph-safe execution kernel [1].

Citations:

• 1: https://github.com/local-inference-lab/b12x
• 2: GitHub pull request 49 in local-inference-lab/b12x (link omitted to avoid creating a cross-reference)

🏁 Script executed:

#!/bin/bash
set -eu
printf '%s\n' '--- tracked dependency and version metadata ---'
git ls-files | grep -Ei '(^|/)(pyproject|setup|requirements|constraints|environment|docker|readme|makefile|lock)' | head -200
printf '%s\n' '--- exact b12x mentions in tracked metadata ---'
git grep -n -i -E 'b12x|dense_mla' -- \
  'pyproject.toml' 'setup.py' 'setup.cfg' 'requirements*' 'constraints*' \
  'Dockerfile*' 'docker/**' '.github/**' 'README*' 2>/dev/null | head -200 || true

Repository: local-inference-lab/vllm

Length of output: 4643

🏁 Script executed:

#!/bin/bash
set -eu
for repo in bhaktatejas922/b12x local-inference-lab/b12x; do
  echo "--- $repo dense MLA paths ---"
  curl -fsSL "https://api.github.com/repos/$repo/git/trees/f82cfc4b?recursive=1" |
    jq -r '.tree[]?.path' | grep -Ei 'dense.*mla|mla.*api|attention/mla' | head -100 || true
done
echo '--- referenced upstream API source ---'
curl -fsSL 'https://raw.githubusercontent.com/bhaktatejas922/b12x/f82cfc4b/b12x/attention/mla/api.py' |
  grep -n -A35 -B10 -E 'def bind|query_cache_seqlens|class .*Binding|class .*Plan' | head -220
echo '--- local-inference-lab API source candidates ---'
for path in b12x/attention/mla/api.py b12x/attention/dense_mla.py b12x/attention/dense_mla/api.py; do
  url="https://raw.githubusercontent.com/local-inference-lab/b12x/f82cfc4b/$path"
  status=$(curl -L -s -o /tmp/b12x-api "$url" -w '%{http_code}')
  if [ "$status" = 200 ]; then
    echo "--- $path ---"
    grep -n -A35 -B10 -E 'def bind|query_cache_seqlens|class .*Binding|class .*Plan' /tmp/b12x-api | head -220
  fi
done

Repository: local-inference-lab/vllm

Length of output: 688

🏁 Script executed:

#!/bin/bash
set -eu
base='https://raw.githubusercontent.com/bhaktatejas922/b12x/f82cfc4b/'
for path in \
  b12x/attention/mla/__init__.py \
  b12x/attention/mla/api.py \
  b12x/attention/mla/workspace.py \
  b12x/attention/mla/kernel.py \
  b12x/attention/mla/kernel_onepass.py \
  b12x/attention/mla/split.py \
  b12x/attention/mla/traits.py; do
  echo "--- $path ---"
  curl -fsSL "$base$path" |
    grep -n -A45 -B12 -E 'def bind|bind[[:space:]]*=|query_cache_seqlens|dense_mla|def plan|class Caps' |
    head -180 || true
done

Repository: local-inference-lab/vllm

Length of output: 4183

🌐 Web query:

"query_cache_seqlens" "dense_mla.bind"

💡 Result:

The terms query_cache_seqlens and dense_mla.bind are associated with the implementation and usage of DeepSeek's FlashMLA library [1][2]. query_cache_seqlens is a required parameter for the get_mla_metadata function in the FlashMLA library [1][3]. It represents the sequence lengths per batch (typically provided as an int32 torch.Tensor with shape (batch_size,)) which the library uses to calculate the necessary tile scheduler metadata for MLA (Multi-Head Latent Attention) decoding [1][3]. Regarding dense_mla.bind, this does not appear to be a standard or public API in the official DeepSeek FlashMLA repository [1][2]. The library's core API centers on functions such as get_mla_metadata and flash_mla_with_kvcache [1][4][5]. It is possible that dense_mla.bind refers to an internal, framework-specific, or binding-layer implementation detail (e.g., within an integration like vLLM, Sglang, or a custom Rust binding) that handles the binding of dense MLA kernels [3][6][7]. If you are working with a specific framework integration, please consult that framework's internal source code or specific API documentation, as it may use wrappers or binding logic that exposes such methods [6][7].

Citations:

• 1: https://github.com/deepseek-ai/FlashMLA/blob/main/README.md
• 2: http://www.github.com/deepseek-ai/FlashMLA
• 3: https://leeroopedia.com/index.php/Implementation:Sgl_project_Sglang_Flash_MLA
• 4: https://github.com/baonudesifeizhai/FlashMLA
• 5: https://github.com/PFCCLab/FlashMLA
• 6: https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/ops/flashmla.py
• 7: https://github.com/mixlayer/flashmla-rs

🌐 Web query:

"query_cache_seqlens" b12x

💡 Result:

In the b12x library, cache_seqlens refers to a specific tensor used for per-request KV cache management [1][2]. It is a tensor of shape [batch] that stores the current sequence length (the number of tokens currently stored in the KV cache) for each request in a batch [1][2]. The KVCacheManager component is responsible for building this tensor during the scheduling process [2]. It is essential for attention kernels, such as b12x_paged_attention_forward, which require these lengths to correctly locate and process tokens within the paged KV cache memory [1][2]. Key technical details regarding cache_seqlens: - Purpose: It tracks the occupancy of the KV cache for each request, allowing the engine to correctly handle extend (prefill) and decode operations [1][2]. - Construction: The KVCacheManager.build_cache_seqlens method generates this tensor by aggregating the cache_len property of active requests [1]. - Usage: It is passed alongside other metadata, such as the page_table and cu_seqlens_q, to the model runner and attention kernels to define the boundaries of the KV cache operations [3][4][5]. The term b12x itself refers to a specialized CUDA kernel library (using a CuTe DSL) optimized for NVIDIA Blackwell (SM120/SM121) hardware [6][7]. It provides high-performance implementations for tasks like Paged Attention, NVFP4 quantization, and Mixture of Experts (MoE) inference [2][8][6].

Citations:

• 1: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/serve/cache/kv_cache.py
• 2: https://deepwiki.com/bhaktatejas922/b12x/6.3-kv-cache-and-memory-management
• 3: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/serve/model/attention.py
• 4: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/serve/engine/runner.py
• 5: https://github.com/bhaktatejas922/b12x/blob/f82cfc4b/serve/engine/serving.py
• 6: https://deepwiki.com/bhaktatejas922/b12x/1.1-getting-started-and-configuration
• 7: https://github.com/local-inference-lab/b12x
• 8: https://deepwiki.com/bhaktatejas922/b12x/6.2-model-recipes-and-loading

🏁 Script executed:

#!/bin/bash
set -eu
for repo in bhaktatejas922/b12x local-inference-lab/b12x; do
  echo "--- $repo repository metadata ---"
  curl -fsSL "https://api.github.com/repos/$repo" |
    jq -r '"default_branch=\(.default_branch) archived=\(.archived) pushed_at=\(.pushed_at)"'
  branch=$(curl -fsSL "https://api.github.com/repos/$repo" | jq -r '.default_branch')
  echo "--- $repo tree on $branch ---"
  curl -fsSL "https://api.github.com/repos/$repo/git/trees/$branch?recursive=1" |
    jq -r '.tree[]?.path' | grep -Ei 'dense.*mla|mla.*dense|attention/mla' | head -120 || true
done

Repository: local-inference-lab/vllm

Length of output: 1310

🏁 Script executed:

#!/bin/bash
set -eu
base='https://raw.githubusercontent.com/local-inference-lab/b12x/master/'
for path in \
  b12x/attention/dense_mla/__init__.py \
  b12x/attention/dense_mla/api.py \
  b12x/attention/dense_mla/planner.py \
  b12x/attention/dense_mla/_io.py \
  tests/attention/test_dense_mla.py \
  pyproject.toml; do
  echo "--- $path ---"
  curl -fsSL "$base$path" |
    grep -n -A55 -B15 -E 'def bind|query_cache_seqlens|class .*Binding|class .*Caps|def plan|dense_mla' |
    head -260 || true
done

Repository: local-inference-lab/vllm

Length of output: 18697

🏁 Script executed:

#!/bin/bash
set -eu
base='https://raw.githubusercontent.com/local-inference-lab/b12x/master/'
for path in \
  b12x/attention/dense_mla/_scratch.py \
  b12x/attention/dense_mla/_forward.py \
  b12x/attention/dense_mla/_kernel.py \
  b12x/_version.py \
  pyproject.toml; do
  echo "--- $path ---"
  curl -fsSL "$base$path" |
    grep -n -A80 -B20 -E 'def bind|query_cache_seqlens|class Binding|class Plan|class Caps|version|project|requires-python' |
    head -360 || true
done
echo '--- all query_cache_seqlens references in upstream dense MLA tests/source ---'
for path in \
  b12x/attention/dense_mla/_scratch.py \
  b12x/attention/dense_mla/_forward.py \
  b12x/attention/dense_mla/_kernel.py \
  b12x/attention/dense_mla/_reference.py \
  tests/attention/test_dense_mla.py; do
  echo "--- $path ---"
  curl -fsSL "$base$path" | grep -n -C4 'query_cache_seqlens' || true
done

Repository: local-inference-lab/vllm

Length of output: 13897

         remote_tokens = torch.tensor(tokens, dtype=torch.int64, device=self.device)
+        # -1 is the "no draft" sentinel; anything else must be a real token id.
+        if bool(
+            ((remote_tokens < -1) | (remote_tokens >= self.vocab_size)).any()
+        ):
+            raise ValueError(
+                "Remote DSpark returned a token id outside "
+                f"[-1, {self.vocab_size})"
+            )
         active_gpu = torch.tensor(active_indices, dtype=torch.int64, device=self.device)

Set self.vocab_size = vllm_config.model_config.get_vocab_size() in __init__.